package cn.lysoy.demo.aspect;

import cn.lysoy.demo.annotation.AccessLimit;
import cn.lysoy.demo.constant.LimitTypeEnum;
import cn.lysoy.demo.constant.CommonConstant;
import cn.lysoy.demo.constant.RespBeanEnum;
import cn.lysoy.demo.exception.GlobalException;
import cn.lysoy.demo.utils.IPUtils;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;

/**
 * @author SongYang
 */
@Slf4j
@Aspect
@Component
@AllArgsConstructor(onConstructor = @__(@Autowired))
public class LimitAspect {
    private final RedisTemplate redisTemplate;

    @Pointcut("@annotation(cn.lysoy.demo.annotation.AccessLimit)")
    public void pointcut() {}

    @Around("pointcut()")
    public Object around(ProceedingJoinPoint point) throws Throwable {
        HttpServletRequest request = ((ServletRequestAttributes) Objects.requireNonNull(RequestContextHolder.getRequestAttributes())).getRequest();
        MethodSignature signature = (MethodSignature) point.getSignature();
        Method method = signature.getMethod();
        AccessLimit limitAnnotation = method.getAnnotation(AccessLimit.class);
        LimitTypeEnum limitTypeEnum = limitAnnotation.limitType();
        String name = limitAnnotation.name();
        String key;
        int limitPeriod = limitAnnotation.period();
        int limitCount = limitAnnotation.count();
        String separator = limitAnnotation.separator();
        // 自行添加和修改limitType和key生成规则
        switch (limitTypeEnum) {
            case TOKEN:
                key = limitTypeEnum.getValue()+separator+request.getHeader(CommonConstant.DEFAULT_TOKEN_KEY);
                break;
            case IP:
                key = limitTypeEnum.getValue()+separator+ IPUtils.getIpAddr(request);
                break;
            case KEY:
                key = limitTypeEnum.getValue()+separator+limitAnnotation.key();
                break;
            case METHOD:
                key = limitTypeEnum.getValue()+separator+StringUtils.upperCase(method.getName());
                break;
            default:
                key = limitTypeEnum.getValue()+separator+"default";
        }
        if (limitAnnotation.sessionId()) {
            key = key + separator + request.getRequestedSessionId();
        }
        List<String> keys = new ArrayList<>(1);
        String redisKey = limitAnnotation.prefix() + separator + key;
        keys.add(redisKey);
        String luaScript = buildLuaScript();
        RedisScript<Long> redisScript = new DefaultRedisScript<>(luaScript, Long.class);
        Long count = (Long) redisTemplate.execute(redisScript, keys,limitPeriod, limitCount);
        log.info("第{}次访问key为 {}，描述为 [{}] 的接口", count, keys, name);
        if (count == 1) {
            return point.proceed();
        } else {
            throw new GlobalException(RespBeanEnum.ACCESS_LIMIT_ERROR);
        }
    }

    /**
     * 限流脚本,返回并执行计算器自加
     * 调用的时候不超过阈值，返回1。
     * 否则返回0
     *
     * @return lua脚本
     */
    private String buildLuaScript() {
        return
        "local key1 = KEYS[1]"+
        "\nlocal val = redis.call('incr', key1)"+
        "\nlocal ttl = redis.call('ttl', key1)"+
        "\nlocal expire = ARGV[1]"+
                "\nlocal times = ARGV[2]"+
        "\nif val == 1 then"+
        "\nredis.call('expire', key1, tonumber(expire))"+
        "\nelse"+
        "\nif ttl == -1 then"+
        "\nredis.call('expire', key1, tonumber(expire))"+
        "\nend"+
        "\nend"+
        "\nif val > tonumber(times) then"+
        "\nreturn 0"+
        "\nend"+
        "\nreturn 1";
    }

}
